#!/usr/bin/env python

# TODO , support python3 in future

import sys,os,usb1,array,struct
import collections
def GetInterface(setting):
    """Get the class, subclass, and protocol for the given USB setting."""
    return (setting.getClass(), setting.getSubClass(), setting.getProtocol())
def PrintUsbInformation(device,file=sys.stdout):
    pass
    #~ print('ID %04x:%04x' % (device.getVendorID(), device.getProductID()), file=file)
    #~ print('\tBUS ID', '.'.join(str(x) for x in [device.getBusNumber()] + device.getPortNumberList()),file=file)
    #~ print( '\tDeviceAddress', device.getDeviceAddress())
    #~ try:
        #~ print( '\tSerialNumber',device.getSerialNumber(),file=file)
        #~ print( '\tManufacturer',device.getManufacturer(),file=file)
        #~ print( '\tProduct',device.getProduct()     ,file=file)

        #~ for i in device.iterSettings():
            #~ print ('\tInterfaces',GetInterface(i))
    #~ except usb1.USBErrorAccess as e:
        #~ print ('\tException:',e,file=file)
        #~ import traceback
        #~ x=traceback.format_exc()
        #~ for line in x.splitlines():
            #~ print ('\t',line)



class PermissionDeny(Exception):
    pass

def checkUsbDevice(device,filters=[]):
    access=True
    try:
        sn=device.getSerialNumber()
        manu=device.getManufacturer()
        product=device.getProduct()

    except usb1.USBErrorAccess:
        access=False
    ret=[]
    for fn in filters:
        r=fn(device)
        if r is not None:
            ret+=[r]
    if access is False :
        if len(ret)>0:
            raise PermissionDeny('ID %04x:%04x , no access permission,it match %s ' % (device.getVendorID(), device.getProductID(),str(ret)))
    return ret
def hotplug_loop(context, callback=None):
    def default_handler(context, device, event):
        if event == usb1.HOTPLUG_EVENT_DEVICE_ARRIVED:
            print('Device Arrived')
            PrintUsbInformation(device)
        elif event == usb1.HOTPLUG_EVENT_DEVICE_LEFT:
            print('Device Left',device)
    if callback==None:
        callback=default_handler
    if not context.hasCapability(usb1.CAP_HAS_HOTPLUG):
        print( 'Hotplug support is missing. Please update your libusb version.')
        return
    print ('Registering hotplug callback...')
    opaque = context.hotplugRegisterCallback(callback)
    try:
        while True:
            context.handleEvents()
    except (KeyboardInterrupt, SystemExit):
        raise
devices=[]
def AML_BootRom(device):
    if (device.getVendorID(), device.getProductID()) not in [(0x1b8e,0xc003)]:
        return None
    return 'Amlogic Boot Rom'
def ADB_Devices(device):

    for setting in device.iterSettings():
      if GetInterface(setting) == (0xff,0x42,0x1):
        return 'ADB Device'
    return None
# Bootrom Command consts
DEVICE2HOST=0xc0
HOST2DEVICE=0x40
#command tables
AM_REQ_WRITE_MEM    =0x01
AM_REQ_READ_MEM     =0x02
AM_REQ_FILL_MEM     =0x03
AM_REQ_MODIFY_MEM   =0x04
AM_REQ_RUN_IN_ADDR  =0x05
AM_REQ_WRITE_AUX    =0x06
AM_REQ_READ_AUX     =0x07
AM_REQ_WR_LARGE_MEM =0x11
AM_REQ_RD_LARGE_MEM =0x12
AM_REQ_IDENTIFY_HOST=0x20
#AM_REQ_TPL_CMD      =0x30
#AM_REQ_TPL_STAT     =0x31
import libusb1
class UsbInterface():
    def __init__(self,device,setting):
        self._device=device
        self._handle=self._device.open()
        self._setting=setting
        self._iface_number=setting.getNumber()
        for endpoint in setting.iterEndpoints():
            address = endpoint.getAddress()
            if address & libusb1.USB_ENDPOINT_DIR_MASK:
                self._read_endpoint = endpoint
            else:
                self._write_endpoint = endpoint
    def open(self):
        self._handle.claimInterface(self._iface_number)
        return self
    def close(self):
        self._handle.releaseInterface(self._iface_number)
    def __enter__(self):
        return self.open()

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
    def bulkRead(self,length,timeout=0):
        print(self._read_endpoint.getAddress())
        return self._handle.bulkRead(self._read_endpoint.getAddress(),length,timeout=timeout)
    def bulkWrite(self,data,timeout=0):
        #~ print(data)
        return self._handle.bulkWrite(self._write_endpoint.getAddress(),data,timeout=timeout)

    def __getattr__(self,name):
        if name in ['controlRead','controlWrite']:
            return getattr(self._handle,name)
        return object.__getattr__(self,name)

class AmlBootRom():
    def __init__(self,device):
        self._device=device
        for setting in device.iterSettings():
            if GetInterface(setting) == (0xff,0,0):
                self._setting=setting
    def Open(self):
        return UsbInterface(self._device,self._setting)
    def IdentifyDevice(self):
        with self.Open() as iface:
            ret=iface.controlRead(DEVICE2HOST,AM_REQ_IDENTIFY_HOST,0,0,16)
        return ret

    def _control_rd(self,iface,addr,count=4):
        assert count < 512
        return iface.controlRead(DEVICE2HOST,AM_REQ_READ_MEM,(addr>>16)&0xffff,addr&0xffff,count)
    def _control_wr(self,iface,addr,data):
        assert len(data) < 512
        ret=iface.controlWrite(HOST2DEVICE,AM_REQ_WRITE_MEM,(addr>>16)&0xffff,addr&0xffff,data)
        return ''
    def _bulk_rd(self,iface,addr,count=4096):
        for shift in range(12,8,-1):
            block_len=1<<shift
            if count %block_len < 512:
                break
        num_blocks=count >> shift
        if num_blocks == 0:
            return ''
        cmd_data=struct.pack("IIII",addr,num_blocks*block_len,0,0)
        iface.controlWrite(HOST2DEVICE,AM_REQ_RD_LARGE_MEM,block_len,num_blocks,cmd_data)
        return iface.bulkRead(block_len*num_blocks)
    def _bulk_wr(self,iface,addr,data):
        count=len(data)
        for shift in range(12,8,-1):
            block_len=1<<shift
            if count %block_len < 512:
                break
        num_blocks=count >> shift
        if num_blocks == 0:
            return data
        w_count=num_blocks*block_len
        import struct
        cmd_data=struct.pack("IIII",addr,num_blocks*block_len,0,0)
        iface.controlWrite(HOST2DEVICE,AM_REQ_WR_LARGE_MEM,block_len,num_blocks,cmd_data)
        iface.bulkWrite(data[:w_count])
        return data[w_count:]
    def Rd(self,addr,count=4):
        ret=b''
        with self.Open() as iface:

            while len(ret) < count:
                l_count=count - len(ret)
                if l_count < 512:
                    ret+=self._control_rd(iface,addr+len(ret),count=l_count)
                else:
                    ret+=self._bulk_rd(iface,addr+len(ret),count=l_count)
        return ret
    def Wr(self,addr,data):
        r=data
        total=len(r)
        written=0
        with self.Open() as iface:
            while len(r)>0:
                l_count=len(r)
                written=total-l_count
                print(written,total,l_count)
                if l_count < 512:
                    r=self._control_wr(iface,addr+written,r)
                else:
                    r=self._bulk_wr(iface,addr+written,r)
    def Run(self,addr,data=None):
        if data:
            self.Wr(addr,data)
        with self.Open() as iface:
            cmd_data=struct.pack("I",1)
            iface.controlWrite(HOST2DEVICE,AM_REQ_RUN_IN_ADDR,(addr>>16)&0xfffff,addr&0xffff,cmd_data)
    def readl(self,addr):
        x=self.Rd(addr)
        return array.array('I',x)[0]
    def writel(self,addr,value):
        self.Wr(addr,struct.pack("I",value))


    def __str__(self):
        return 'Amlogic Boot Rom' + repr(self._device)
#~ def test_ReadReg(x):
    #~ meta_rev=x.Rd(0xc8834400 + (0x6a << 2))
    #~ hw_rev=x.Rd((0x1f53 <<2) + 0xc1100000)
    #~ for i in meta_rev + hw_rev:
        #~ print (hex(i))
#~ def test_identifyhost(x):
    #~ print (repr(x.IdentifyDevice()))
#~ def test_writeReg(x):
    #~ #test Write memory
    #~ import array
    #~ for i in array.array('I',x.Rd(0xd9000000,8)):
        #~ print('Before',hex(i))
    #~ x.Wr(0xd9000000,array.array('I',[0x1234,0x5678,0x9abc,0xdef0]).tostring())
    #~ for i in array.array('I',x.Rd(0xd9000000,8)):
        #~ print('After',hex(i))
def test_bulk_write(x):
    count=513
    addr=0xd9000000
    data=array.array('I',[(i&0x3f)+0x20 for i in range(count)]).tobytes()
    x.Wr(addr,data)
    s=x.Rd(0xd9000000,len(data))
    y=array.array('I',s)
    i=0
    print('%08x %08x %08x %08x' %(y[i],y[i+1],y[i+2],y[i+3]))
    if s != data :
        print("FAIL",len(s),len(data),type(s),type(data))
    else:
        print("PASS")
def run_all_tests(x):
    tbl=globals().copy()
    for i in tbl:
        if i.startswith('test_') and isinstance(tbl[i], collections.Callable):
            print('='*8,i,'='*8)
            tbl[i](x)
            print('='*8,i,'END','='*8)

def test_spi_read(x):
    print(type(array.array('I',[0x2aa949]).tobytes()))
    print(type(array.array('I',[0x2aa949]).tostring()))
    x.Wr(0xc1108c88,array.array('I',[0x2aa949]).tobytes())
    pinmux45=array.array('I',x.Rd(0xc8834400 + (0x30 << 2),count=8))
    pinmux45[0]&= ~((1<<31) | (7<<22) | (1<<20))
    #~ *P_PERIPHS_PIN_MUX_5 |= 0xf;
    pinmux45[1]|= 0xf
    for i in pinmux45:
        print(hex(i))
    x.Wr(0xc8834400 + (0x30 << 2),pinmux45.tostring())
    y=array.array('I',x.Rd(0xcc000000+0x10000,8192))
    i=0
    print('%08x %08x %08x %08x' %(y[i],y[i+1],y[i+2],y[i+3]))
def mytest_run(x):
    print('%08x %08x %08x '% (x.readl(0xc8834430),x.readl(0xc8834434),x.readl(0xc8834438)) )
    for i in range(8):
        print(hex(0xc8834400 + ((0x2c + i)<< 2)),hex(x.readl(0xc8834400 + ((0x2c + i)<< 2))))
    with open('spi_spl.spl','rb') as f:
        x.Run(0xd9000000,f.read())
def filter_know_devices(context, device, event):
    if event == usb1.HOTPLUG_EVENT_DEVICE_ARRIVED:
        PrintUsbInformation(device)
        dev=checkUsbDevice(device,filters=[AML_BootRom,ADB_Devices])
        if len(dev)>0:
            print(dev)
            devices.append(device)

            x=AmlBootRom(device)
            run_all_tests(x)
            #~ mytest_run(x)

            #~ test_writeReg(x)
            test_bulk_write(x)

    elif event == usb1.HOTPLUG_EVENT_DEVICE_LEFT:
        print(device,'left')
        if device in devices:
            devices.remove(device)
def main():
    with usb1.USBContext() as context:
        hotplug_loop(context,callback=filter_know_devices)
